# sample_predictor.py
import numpy as np
from asm_feature_extractor import process_asm_file

def predict_sample(model_path, asm_path, custom):
    """
    使用训练好的模型预测样本，返回预测结果和概率分布
    
    Args:
        model_path: BERT模型路径
        asm_path: 汇编文件路径
        custom: 自定义模型对象 (joblib.load 返回的 pipeline)
    
    Returns:
        prediction: 预测结果（与训练时标签一致的原始标签）
        probability: 完整的概率分布数组
    """
    # 处理ASM文件
    feature_vector = process_asm_file(model_path, asm_path)
    
    if feature_vector is None:
        print("无法从文件中提取特征")
        return None, None
    
    feature_vector = np.array(feature_vector).reshape(1, -1)
    
    # 预测原始结果（0-8）
    raw_prediction = custom.predict(feature_vector)[0]
    # 将预测结果还原为原始标签（1-9）
    prediction = raw_prediction + 1

    # 预测
    if hasattr(custom, 'predict_proba'):
        probability = custom.predict_proba(feature_vector)[0]
    else:
        # 对于不支持概率预测的模型，使用决策函数并 softmax
        decision_values = custom.decision_function(feature_vector)[0]
        exps = np.exp(decision_values - np.max(decision_values))
        probability = exps / np.sum(exps)
    
    return prediction, probability
